package task

import (
	"context"
	"fmt"
	"sync"
	proto "xagent/proto/v1"
)

// Priority represents task priority levels
type Priority int

const (
	PriorityLow Priority = iota
	PriorityMedium
	PriorityHigh
)

// Scheduler manages task scheduling and execution
type Scheduler struct {
	tasks    map[string]*proto.Task
	queue    []*proto.Task
	mutex    sync.RWMutex
	maxTasks int
}

// NewScheduler creates a new task scheduler
func NewScheduler() *Scheduler {
	return &Scheduler{
		tasks:    make(map[string]*proto.Task),
		queue:    make([]*proto.Task, 0),
		maxTasks: 10,
	}
}

// ScheduleTask adds a task to the scheduler
func (s *Scheduler) ScheduleTask(ctx context.Context, task *proto.Task) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	if task.Metadata == nil {
		task.Metadata = make(map[string]string)
	}

	// Default to medium priority if not specified
	priority := "medium"
	if p, ok := task.Metadata["priority"]; ok {
		priority = p
	}
	task.Metadata["priority"] = priority

	s.tasks[task.Id] = task
	s.queue = append(s.queue, task)
	return nil
}

// String returns the string representation of Priority
func (p Priority) String() string {
	switch p {
	case PriorityLow:
		return "low"
	case PriorityMedium:
		return "medium"
	case PriorityHigh:
		return "high"
	default:
		return "unknown"
	}
}

// GetNextTask gets the next task to execute
func (s *Scheduler) GetNextTask() *proto.Task {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	if len(s.queue) == 0 {
		return nil
	}

	task := s.queue[0]
	s.queue = s.queue[1:]
	delete(s.tasks, task.Id)
	return task
}

// HasTasks checks if there are any tasks in the queue
func (s *Scheduler) HasTasks() bool {
	s.mutex.RLock()
	defer s.mutex.RUnlock()
	return len(s.queue) > 0
}

// GetTaskCount returns the number of tasks in the queue
func (s *Scheduler) GetTaskCount() int {
	s.mutex.RLock()
	defer s.mutex.RUnlock()
	return len(s.queue)
}

// GetTask gets a task by ID
func (s *Scheduler) GetTask(id string) (*proto.Task, bool) {
	s.mutex.RLock()
	defer s.mutex.RUnlock()
	task, ok := s.tasks[id]
	return task, ok
}

// RemoveTask removes a task from the scheduler
func (s *Scheduler) RemoveTask(id string) {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	delete(s.tasks, id)
	for i, task := range s.queue {
		if task.Id == id {
			s.queue = append(s.queue[:i], s.queue[i+1:]...)
			break
		}
	}
}

// CancelTask cancels a task by ID
func (s *Scheduler) CancelTask(ctx context.Context, taskID string) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	if _, exists := s.tasks[taskID]; !exists {
		return fmt.Errorf("task with ID %s not found", taskID)
	}

	// Remove from tasks map
	delete(s.tasks, taskID)

	// Remove from queue
	for i, task := range s.queue {
		if task.Id == taskID {
			s.queue = append(s.queue[:i], s.queue[i+1:]...)
			break
		}
	}

	return nil
}

// GetAvailableResources returns a map of available resources
func (s *Scheduler) GetAvailableResources() map[string]float64 {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	// Calculate available resources based on current task load
	maxTasks := float64(s.maxTasks)
	currentTasks := float64(len(s.queue))

	return map[string]float64{
		"capacity": maxTasks - currentTasks,
		"memory":   100.0, // Default memory availability percentage
		"cpu":      100.0, // Default CPU availability percentage
	}
}

// GetMaxConcurrent returns the maximum number of concurrent tasks
func (s *Scheduler) GetMaxConcurrent() int {
	s.mutex.RLock()
	defer s.mutex.RUnlock()
	return s.maxTasks
}

// GetMetrics returns scheduler metrics
func (s *Scheduler) GetMetrics() interface{} {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	return map[string]interface{}{
		"total_tasks":          len(s.tasks),
		"queued_tasks":         len(s.queue),
		"max_concurrent_tasks": s.maxTasks,
		"queue_utilization":    float64(len(s.queue)) / float64(s.maxTasks),
	}
}

// GetPendingTasks returns all tasks that are pending execution
func (s *Scheduler) GetPendingTasks(ctx context.Context) ([]*proto.Task, error) {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	pendingTasks := make([]*proto.Task, len(s.queue))
	copy(pendingTasks, s.queue)
	return pendingTasks, nil
}

// GetResourceUsage returns the current resource usage
func (s *Scheduler) GetResourceUsage() map[string]float64 {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	maxTasks := float64(s.maxTasks)
	currentTasks := float64(len(s.queue))

	return map[string]float64{
		"capacity": currentTasks / maxTasks * 100.0,
		"memory":   0.0, // Default memory usage percentage
		"cpu":      0.0, // Default CPU usage percentage
	}
}

// GetTaskQueue returns the current number of tasks in the queue
func (s *Scheduler) GetTaskQueue() int {
	s.mutex.RLock()
	defer s.mutex.RUnlock()
	return len(s.queue)
}

// GetTaskStatus returns the status of a task
func (s *Scheduler) GetTaskStatus(ctx context.Context, taskID string) (proto.TaskStatus, error) {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	task, exists := s.tasks[taskID]
	if !exists {
		return proto.TaskStatus_TASK_STATUS_UNSPECIFIED, fmt.Errorf("task with ID %s not found", taskID)
	}

	return task.Status, nil
}
